# source.py: Sends a UDP packet once per second to a destination
import sys
import threading
import time
import socket
import math
from optparse import OptionParser
import csv

class Source():
  def __init__(self, dest_addr, dest_port, local_port, interval, decay, dest_port2, probe, switch):
    self.dest = (dest_addr, dest_port)
    self.dest2 = (dest_addr, dest_port2)
    self.local = ('0.0.0.0', local_port)
    self.csv_name = 'Source_%s.csv' % time.strftime('%H%M%S')
    try:
      self.listen_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
      self.listen_sock.bind(self.local)
    except socket.error:
      self.local = ('0.0.0.0', local_port - 2000)
      self.listen_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
      self.listen_sock.bind(self.local)
      self.csv_name = 'Source2_%s.csv' % time.strftime('%H%M%S')
    print 'Listening port bound to %d' % self.local[1]
    self.outgoing_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    self.listen_thread = None
    self.current_index = 0
    self.packets = {}  # Store outstanding packets.  Key = packet # (str), Value = Time (float)
    self.probe_packets = {}
    self.interval = interval
    self.decay = decay
    self.average_rtt = 10.0
    self.print_lock = threading.Lock()
    self.packets_lock = threading.Lock()
    self.probe = probe
    self.send_count = 0
    self.probe_count = 0
    self.switch = switch
    self.average_rtt2 = 10.0
    self.csv_file = None
    self.csv_writer = csv.writer(open(self.csv_name, 'w'))
    self.csv_writer.writerow([
      'Packet Type', 'Time Stamp', 'Sequence Number', 'Packet Data', 
      'Destination', 'RTT', 'Average RTT', 'Timeout?'
    ])
    
    # Check that decay is within the range
    if decay > 1.0: decay = 1.0
    elif decay < 0.0: decay = 0.0  
    
  def _listen_runner(self):
    while True:
      # Listen for a new packet
      data, addr = self.listen_sock.recvfrom(1024)
      is_probe_packet = False
      
      # Retrieve the time this packet was sent to calculate RTT
      with self.packets_lock:
        if data in self.packets:
          start_time = self.packets[data]
          del self.packets[data]
        elif data in self.probe_packets:
          start_time = self.probe_packets[data]
          del self.probe_packets[data]
          is_probe_packet = True
        else:
          start_time = None
      elapsed_time = time.time() - start_time if start_time else None
      
      # Recalculate A_n = (1-q)*A_{n-1} + q*T_n, only if packet didn't time out
      if not is_probe_packet:
        if elapsed_time < 4 * self.average_rtt:
          # Note: we made elapsed time a minimum of 1ms, otherwise weird things happen
          self.average_rtt = (1 - self.decay) * self.average_rtt + self.decay * (elapsed_time if elapsed_time > 0.0 else 0.001)
          timeout = ''
        else:
          timeout = 'Timeout'
      else:
        if elapsed_time < 4 * self.average_rtt2:
          # Note: we made elapsed time a minimum of 1ms, otherwise weird things happen
          self.average_rtt2 = (1 - self.decay) * self.average_rtt2 + self.decay * (elapsed_time if elapsed_time > 0.0 else 0.001)
          timeout = ''
          self.probe_count += 1
        else:
          timeout = 'Timeout'
      
      # Print out that we successfully received the packet
      packet_type = 'ack' if not is_probe_packet else 'p_ack'
      with self.print_lock:
        if elapsed_time is not None:
          print '%s\t%f\t%s\t%f\t%s' % (packet_type, time.time(), data, elapsed_time, timeout)
          self.csv_writer.writerow([packet_type, time.time(), data, data, 'N/A', 
              elapsed_time, self.average_rtt if not is_probe_packet else self.average_rtt2, 
              timeout])
        else:
          print '%s\t%f\t%s\t%s\t%s' % (packet_type, time.time(), data, '-1', timeout)
        if not is_probe_packet:
          print 'A_n\t%f' % self.average_rtt
        else:
          print 'A_n2\t%f' % self.average_rtt2
    
  def run(self):
    """Runs the main program"""
    
    # Start listening
    print 'Start the listener...'
    self.listen_thread = threading.Thread(target=self._listen_runner)
    self.listen_thread.daemon = True
    self.listen_thread.start()
    
    # Loop forever
    while True:
      # Check to see if we want to switch queues
      if self.probe_count > self.switch:
        if self.average_rtt2 < self.average_rtt and (math.fabs(self.average_rtt - self.average_rtt2) > (self.average_rtt*.1)):
          # Switch queues and reset counters
          self.dest, self.dest2 = self.dest2, self.dest
          self.average_rtt, self.average_rtt2 = self.average_rtt2, self.average_rtt
          self.probe_count = 0
          self.send_count = 0
          with self.print_lock:
            print 'Switched! Using %s now' % str(self.dest)
            self.csv_writer.writerow(['switch', time.time(), str(self.dest), 
                'N/A', self.average_rtt, 'N/A'])
      host, port = self.listen_sock.getsockname()
      msg = '%s %s:%s' % (str(self.current_index), host, port)
      data = str(self.current_index)
      self.current_index += 1
      with self.packets_lock:
        self.outgoing_sock.sendto(msg, self.dest)
        self.packets[data] = time.time()
      with self.print_lock:
        print 'send\t%f\t%s\tto %s' % (time.time(), msg, self.dest)
        self.csv_writer.writerow(['send', time.time(), data, msg, str(self.dest), 'N/A', 
                                   self.average_rtt, 'N/A'])
      self.send_count += 1
      
      # Send a probe packet
      if self.send_count >= self.probe:
        self.send_count = 0
        host, port = self.listen_sock.getsockname()
        msg = '%s %s:%s' % (str(self.current_index), host, port)
        data = str(self.current_index)
        self.current_index += 1
        with self.packets_lock:
          self.outgoing_sock.sendto(msg, self.dest2)
          self.probe_packets[data] = time.time()
        with self.print_lock:
          print 'probe\t%f\t%s\tto %s' % (time.time(), msg, self.dest2)
          self.csv_writer.writerow(['probe', time.time(), data, msg, str(self.dest2), 'N/A', 
                                     self.average_rtt2, 'N/A'])
      # Wait for T second(s)
      time.sleep(self.interval)
    
    

if __name__ == '__main__':
  # Parse arguments [dest_addr=localhost] [dest_port=9001] [local_port=9003] [T=1]
  parser = OptionParser()
  parser.add_option("-d", "--dest_addr", dest="dest_addr", default='127.0.0.1',
                    help="Destination IP address")
  parser.add_option("-p", "--dest_port", dest="dest_port", default=9001, type='int',
                    help="Destination port number")
  parser.add_option("-s", "--dest_port2", dest="dest_port2", default=7001, type='int',
                    help="Destination port number for second queue")
  parser.add_option("-l", "--local_port", dest="local_port", default=9003, type='int',
                    help="Local (listening) port number")
  parser.add_option("-T", "--time_interval", dest="interval", default=1, type='float',
                    help="Time interval between packets (in seconds)")
  parser.add_option("-q", "--decay_constant", dest="decay", default=0.5, type='float',
                    help="Decay constant q for A_n algorithm")
  parser.add_option("-n", "--probe_period", dest="probe", default=3, type='int',
                    help="Number of sent packets between probes")
  parser.add_option("-w", "--switchover_count", dest="switch", default=3, type='int',
                    help="Number of probes to use before deciding to switch queues.")
  (options, args) = parser.parse_args()
  
  # Create a new "Source" object
  print 'Binding sockets...'
  source = Source(options.dest_addr, options.dest_port, options.local_port, 
                  options.interval, options.decay, options.dest_port2, 
                  options.probe, options.switch)
  
  # Run forever
  source.run()
